import os
import sys
import logging
import time
import numpy as np
from onnxruntime import InferenceSession
from transformers import ElectraTokenizer

from .. import helper
from ..onnx_base import ONNXBase


class ONNXInfer(ONNXBase):
    def __init__(self, model_dir=None) -> None:
        super().__init__(model_dir)
        self.seq_len = 384
        self.tokenizer = ElectraTokenizer.from_pretrained(self.model_dir)

        onnx_path = f"{model_dir}/electra.onnx"
        self.onnx_sess = InferenceSession(onnx_path, providers=self.providers)
        self.inputs_name_list = self._get_onnx_input_name(self.onnx_sess)
        self.outputs_name_list = self._get_onnx_output_name(self.onnx_sess)

    def _preproc(self, data):
        tokens = self.tokenizer(data, padding="max_length", max_length=self.seq_len, return_tensors="np")
        input_ids = tokens.input_ids
        attention_mask = tokens.attention_mask
        token_type_ids = tokens.token_type_ids

        logging.info(f"ids shape:{input_ids.shape} {input_ids.dtype}")
        logging.info(f"mask shape:{attention_mask.shape} {attention_mask.dtype}")
        logging.info(f"token shape:{token_type_ids.shape} {token_type_ids.dtype}")

        return input_ids, attention_mask, token_type_ids

    def infer(self, contents):
        input_ids, attention_mask, token_type_ids = self._preproc(contents)

        time_start = time.time()
        input_ids = input_ids.astype(np.int32)
        attention_mask = attention_mask.astype(np.int32)
        token_type_ids = token_type_ids.astype(np.int32)
        start_logits, end_logits = self.onnx_sess.run(self.outputs_name_list, 
                                            {self.inputs_name_list[0]:input_ids, 
                                             self.inputs_name_list[1]:attention_mask,
                                             self.inputs_name_list[2]:token_type_ids})
        infer_time = (time.time() - time_start) * 1000
        logging.debug(f"shape start logits:{start_logits.shape}, end logits:{end_logits.shape}")
        logging.info(f"infer time:{infer_time:.3f} (ms)")

        return start_logits, end_logits